#ifndef __M3N_STATS_H__
#define __M3N_STATS_H__
/*********************************************************************
 * Software License Agreement (Modified BSD License)
 *
 * Copyright (c) 2010, Daniel Munoz
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *   * Neither the name of the copyright holders' organizations nor the
 *     names of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *********************************************************************/

#include <vector>
#include <map>
#include <set>

#include <m3n/random_field.h>

// no error checks
class M3NStats
{
  public:
    static void computeConfusionMatrix(const std::map<unsigned int, RandomField::Node*>& rf_nodes,
                                       const std::vector<unsigned int>& training_labels,
                                       const std::map<unsigned int, unsigned int>& infer_labels,
                                       std::vector<std::vector<unsigned int> >& confusion_matrix)
    {
      unsigned int nbr_labels = training_labels.size();

      // init if necessary
      if (confusion_matrix.size() == 0)
      {
        std::vector<unsigned int> empty_row(nbr_labels, 0);
        confusion_matrix.assign(nbr_labels, empty_row);
      }

      std::map<unsigned int, unsigned int> label2idx;
      for (unsigned int i = 0 ; i < nbr_labels ; i++)
      {
        label2idx[training_labels[i]] = i;
      }

      std::map<unsigned int, unsigned int>::const_iterator iter_infer_labels;
      for (iter_infer_labels = infer_labels.begin(); iter_infer_labels != infer_labels.end() ; iter_infer_labels++)
      {
        unsigned int node_id = iter_infer_labels->first;
        unsigned int infer_label = iter_infer_labels->second;

        const RandomField::Node* node = rf_nodes.find(node_id)->second;
        unsigned int gt_label = node->getLabel();

        confusion_matrix[label2idx[gt_label]][label2idx[infer_label]]++;
      }
    }

    static void computeStats(const std::vector<std::vector<unsigned int> >& confusion_matrix,
                             double& accuracy,
                             std::vector<double>& precisions,
                             std::vector<double>& recalls,
                             std::vector<double>& f1_scores,
                             double& macro_precision,
                             double& macro_recall,
                             double& macro_f1)
    {
      computeAccuracy(confusion_matrix, accuracy);
      computePrecisions(confusion_matrix, precisions);
      computeRecalls(confusion_matrix, recalls);
      computeF1Scores(precisions, recalls, f1_scores);
      computeMeans(precisions, recalls, f1_scores, macro_precision, macro_recall, macro_f1);
    }

    static void computeAccuracy(const std::vector<std::vector<unsigned int> >& confusion_matrix,
                                double& accuracy)
    {
      accuracy = 0.0;
      double total_sum = 0.0;
      for (unsigned int i = 0 ; i < confusion_matrix.size() ; i++)
      {
        for (unsigned int j = 0 ; j < confusion_matrix.size() ; j++)
        {
          total_sum += confusion_matrix[i][j];
        }
        accuracy += confusion_matrix[i][i];
      }

      accuracy /= total_sum;
    }

    static void computePrecisions(const std::vector<std::vector<unsigned int> >& confusion_matrix,
                                  std::vector<double>& precisions)
    {
      unsigned int nbr_labels = confusion_matrix.size();
      precisions.assign(nbr_labels, 0.0);
      for (unsigned int i = 0 ; i < nbr_labels ; i++)
      {
        unsigned int col_sum = 0;
        for (unsigned int j = 0 ; j < nbr_labels ; j++)
        {
          col_sum += confusion_matrix[j][i];
        }

        if (col_sum == 0)
        {
          precisions[i] = 1.0;
        }
        else
        {
          precisions[i] = confusion_matrix[i][i] / static_cast<double> (col_sum);
        }
      }
    }

    static void computeRecalls(const std::vector<std::vector<unsigned int> >& confusion_matrix,
                               std::vector<double>& recalls)
    {
      unsigned int nbr_labels = confusion_matrix.size();
      recalls.assign(nbr_labels, 0.0);
      for (unsigned int i = 0 ; i < nbr_labels ; i++)
      {
        unsigned int row_sum = 0;
        for (unsigned int j = 0 ; j < nbr_labels ; j++)
        {
          row_sum += confusion_matrix[i][j];
        }

        if (row_sum == 0)
        {
          recalls[i] = 1.0;
        }
        else
        {
          recalls[i] = confusion_matrix[i][i] / static_cast<double> (row_sum);
        }
      }
    }

    static void computeF1Scores(const std::vector<double>& precisions,
                                const std::vector<double>& recalls,
                                std::vector<double>& f1_scores)
    {
      unsigned int nbr_labels = precisions.size();
      f1_scores.assign(nbr_labels, 0.0);
      for (unsigned int i = 0 ; i < nbr_labels ; i++)
      {
        f1_scores[i] = (2 * precisions[i] * recalls[i]) / (precisions[i] + recalls[i]);
      }
    }

    static void computeMeans(std::vector<double>& precisions,
                             std::vector<double>& recalls,
                             std::vector<double>& f1_scores,
                             double& macro_precision,
                             double& macro_recall,
                             double& macro_f1)
    {
      macro_precision = 0.0;
      macro_recall = 0.0;
      macro_f1 = 0.0;
      for (unsigned int i = 0 ; i < precisions.size() ; i++)
      {
        macro_precision += precisions[i];
        macro_recall += recalls[i];
        macro_f1 += f1_scores[i];
      }
      macro_precision /= static_cast<double> (precisions.size());
      macro_recall /= static_cast<double> (precisions.size());
      macro_f1 /= static_cast<double> (precisions.size());
    }
};

#endif
